import json
import os.path
import cv2
import tqdm

scenes = ['', 'frame_20190829091111_x_', 'frame_20190905091750_x_', 'frame_20190905103112_x_',
          'frame_20190905111947_x_', 'frame_20190905112522_x_', 'frame_20190905142119_x_',
          'frame_20190905143505_x_', 'frame_20190906150731_x_']


# select_scene的值应该为1-8，分别表示AuAir数据中8个不同的场景
def coco2yoloForAuAir(select_scene):
    scene = scenes[select_scene]
    dataset_root_dir = "../datasets/AuAir/"
    save_path = dataset_root_dir + "scene" + str(select_scene) + "/annotation/"
    annotationPath = dataset_root_dir + "origin/annotations.json"
    data = json.load(open(annotationPath, 'r'))
    annotations = data['annotations']
    categories = data['categories']

    # 将coco格式的数据集转化为yolo格式的数据集
    for cur in annotations:
        name = cur['image_name'].split('.')[0] + '.txt'
        if scene in name:
            txt_path = os.path.join(save_path, name)
            image_path = dataset_root_dir + 'origin/all_images/' + cur['image_name']
            with open(txt_path, 'a', ) as file:
                bbox = cur['bbox']
                for item in bbox:
                    cls = categories[item['class']].lower()
                    cls = 'person' if cls == 'human' else cls
                    top = item['top']
                    left = item['left']
                    width = item['width']
                    height = item['height']
                    line = "{} {} {} {} {}\n".format(cls, left, top, width, height)
                    file.write(line)


def draw_rect(img_path, txt_path, isPredictResult=False):
    """
        将txt中的坐标所表示的bounding box在图像上画出
        img_path 为图片路径
        txt_path 为txt文件路径
        isPredictResult 表示是否是预测框，False 的时候为真值框
    """
    img = cv2.imread(img_path)
    lines = open(txt_path, 'r').readlines()
    for bbox in lines:
        info = bbox.strip().split(' ')
        # 预测框中有一个置信度，在画bounding box 中不需要，所以需要删除掉
        if isPredictResult:
            del info[1]
        cls = info[0]
        left = int(info[1])
        top = int(info[2])
        width = int(info[3])
        height = int(info[4])
        cv2.rectangle(img, (left, top), (left + width, top + height), (0, 255, 0), 2)
    cv2.imshow('img', img)
    cv2.waitKey(0)


def xywh2xcycwh(select_scene):
    scene = scenes[select_scene]
    dataset_root_dir = "../datasets/AuAir/"
    save_path = "../datasets/scene" + str(select_scene) + "/labels/"
    annotationPath = dataset_root_dir + "origin/annotations.json"
    data = json.load(open(annotationPath, 'r'))
    annotations = data['annotations']
    categories = data['categories']

    # 将coco格式的数据集转化为yolo格式的数据集
    for cur in annotations:
        name = cur['image_name'].split('.')[0] + '.txt'
        if scene in name:
            txt_path = os.path.join(save_path, name)
            with open(txt_path, 'a', ) as file:
                bbox = cur['bbox']
                for item in bbox:
                    cls = item['class']
                    top = float(item['top'])
                    left = float(item['left'])
                    width = float(item['width'])
                    height = float(item['height'])
                    x_center = round((left + width/2)/1920, 6)
                    y_center = round((top + height / 2) / 1080, 6)
                    width = round(width / 1920, 6)
                    height = round(height / 1080, 6)
                    line = "{} {} {} {} {}\n".format(cls, x_center, y_center, width, height)
                    file.write(line)


xywh2xcycwh(8)
